import pandas
from elasticsearch import Elasticsearch

from conf import es_url


def get_teacher_info(talker_name, teacher_csv):
    df = pandas.read_csv(teacher_csv)
    teacher = df[(df['code'] == talker_name)]
    if not teacher.empty:
        teacher_info = teacher.iloc[0].to_dict()
    else:
        teacher_info = None
    print(teacher_info)
    return teacher_info


def update_teacher_info(es, index_name, code, teacher_csv, info):
    chat_ids = info[info['code'] == code]['chat_id'].tolist()
    teacher_info = get_teacher_info(code, teacher_csv)
    teacher_info['code'] = 'Mnt1'

    for chat_id in chat_ids:
        update_query = {
            "script": {
                "source": "ctx._source.teacher_info = params.new_teacher_info",
                "params": {
                    "new_teacher_info": teacher_info
                }
            },
            "query": {
                "term": {
                    "chat_id": chat_id
                }
            }
        }

        result = es.update_by_query(index=index_name, body=update_query)

        print(f"成功更新了{result['updated']}个文档，chat_id为: {chat_id}")


if __name__ == "__main__":
    # 创建Elasticsearch客户端
    es = Elasticsearch(es_url)
    # 设置索引名称
    index_name = 'conversation'
    # CSV文件路径
    teacher_csv = 'teacher.csv'
    code = 'Mnt1_W'
    info = pandas.read_csv('andrew.csv')
    update_teacher_info(es, index_name, code, teacher_csv, info)
